from py2neo import Graph, Node, Relationship, Transaction
import json
import os
from typing import List, Dict, Any
import logging
import time
import sys


class Neo4j菜谱导入器:
    def __init__(self, uri: str, 用户名: str, 密码: str, 批量大小: int = 100):
        """初始化导入器

        Args:
            uri: Neo4j数据库连接URI
            用户名: 数据库用户名
            密码: 数据库密码
            批量大小: 每次批量处理的记录数
        """
        self.图数据库 = Graph(uri, auth=(用户名, 密码))
        self.批量大小 = 批量大小
        self._配置日志()
        self.日志.info("初始化Neo4j菜谱导入器")

    def _配置日志(self):
        """配置日志系统"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s [%(levelname)s] %(message)s',
            handlers=[logging.StreamHandler()]
        )
        self.日志 = logging.getLogger(__name__)

    def 深度清理数据库(self):
        """深度清理数据库，包括约束、索引和数据"""
        try:
            开始时间 = time.time()
            self.日志.info("开始深度清理数据库...")

            # 删除所有约束
            约束列表 = self.图数据库.run("SHOW CONSTRAINTS").data()
            for 约束 in 约束列表:
                self.图数据库.run(f"DROP CONSTRAINT {约束['name']}")
                self.日志.debug(f"已删除约束: {约束['name']}")

            # 删除所有索引
            索引列表 = self.图数据库.run("SHOW INDEXES").data()
            for 索引 in 索引列表:
                if 索引['type'] != 'LOOKUP':
                    self.图数据库.run(f"DROP INDEX {索引['name']}")
                    self.日志.debug(f"已删除索引: {索引['name']}")

            # 删除所有数据
            self.图数据库.run("MATCH (n) DETACH DELETE n")

            耗时 = time.time() - 开始时间
            self.日志.info(f"数据库已深度清空 | 耗时: {耗时:.2f}秒")

            # 重新设置约束
            self._设置约束()

        except Exception as 异常:
            self.日志.error(f"深度清空数据库失败: {str(异常)}")
            raise

    def _设置约束(self):
        """设置数据库约束"""
        约束列表 = [
            "CREATE CONSTRAINT IF NOT EXISTS FOR (菜谱:菜谱) REQUIRE 菜谱.名称 IS UNIQUE",
            "CREATE CONSTRAINT IF NOT EXISTS FOR (食材:食材) REQUIRE 食材.名称 IS UNIQUE",
            "CREATE CONSTRAINT IF NOT EXISTS FOR (功效:功效) REQUIRE 功效.名称 IS UNIQUE",
            "CREATE CONSTRAINT IF NOT EXISTS FOR (适宜人群:适宜人群) REQUIRE 适宜人群.描述 IS UNIQUE",
            "CREATE CONSTRAINT IF NOT EXISTS FOR (禁忌人群:禁忌人群) REQUIRE 禁忌人群.描述 IS UNIQUE"
        ]

        for 约束 in 约束列表:
            try:
                self.图数据库.run(约束)
                标签 = 约束.split('FOR')[1].split('REQUIRE')[0].strip()
                self.日志.debug(f"成功创建约束: {标签}")
            except Exception as 异常:
                self.日志.warning(f"创建约束失败: {约束} - {str(异常)}")

    def _显示进度(self, 当前行数: int, 总行数: int, 开始时间: float):
        """显示进度条

        Args:
            当前行数: 已处理的行数
            总行数: 总行数
            开始时间: 开始处理的时间
        """
        耗时 = time.time() - 开始时间
        百分比 = (当前行数 / 总行数) * 100
        旋转图标 = ['-', '\\', '|', '/'][当前行数 % 4]

        # 计算预计剩余时间
        预计剩余时间 = f"预计剩余: {(耗时 / 当前行数 * (总行数 - 当前行数)):.1f}秒" if 当前行数 > 0 else "预计剩余: --"

        # 构建进度条
        进度条长度 = 50
        填充长度 = int(百分比 / 2)
        进度条 = f"[{'=' * 填充长度}{' ' * (进度条长度 - 填充长度)}]"

        # 输出进度信息
        sys.stdout.write(
            f"\r{旋转图标} 进度: {当前行数}/{总行数} {百分比:.1f}% {进度条} "
            f"[{耗时:.1f}秒, {预计剩余时间}]"
        )
        sys.stdout.flush()

    def 导入JSON数据(self, 文件路径: str):
        """导入JSONL格式的菜谱数据

        Args:
            文件路径: JSONL文件路径
        """
        if not os.path.exists(文件路径):
            raise FileNotFoundError(f"文件不存在: {文件路径}")

        self.日志.info(f"开始导入文件: {文件路径}")
        self.深度清理数据库()

        # 计算总行数
        with open(文件路径, 'r', encoding='utf-8') as 文件:
            总行数 = sum(1 for _ in 文件)

        已处理 = 0
        批次 = []
        开始时间 = time.time()

        try:
            with open(文件路径, 'r', encoding='utf-8') as 文件:
                for 行 in 文件:
                    try:
                        数据 = json.loads(行.strip())
                        批次.append(数据)
                        已处理 += 1

                        # 更新进度条
                        self._显示进度(已处理, 总行数, 开始时间)

                        # 批量处理
                        if len(批次) >= self.批量大小:
                            self._处理批次(批次)
                            批次 = []

                    except json.JSONDecodeError as 异常:
                        self.日志.error(f"\nJSON解析错误: {str(异常)}\n问题行: {行[:200]}...")
                    except Exception as 异常:
                        self.日志.error(f"\n处理行时出错: {str(异常)}\n问题行: {行[:200]}...")

                # 处理最后一批数据
                if 批次:
                    self._处理批次(批次)

        finally:
            # 完成进度条
            sys.stdout.write("\n")
            sys.stdout.flush()

            # 输出统计信息
            耗时 = time.time() - 开始时间
            self.日志.info(
                f"导入完成! 共处理 {已处理}/{总行数} 行, "
                f"耗时 {耗时:.2f} 秒, "
                f"平均速度: {已处理 / 耗时:.1f} 行/秒"
            )

    def _处理批次(self, 批次: List[Dict[str, Any]]):
        """处理一批数据

        Args:
            批次: 一批菜谱数据
        """
        事务 = self.图数据库.begin()
        try:
            for 数据 in 批次:
                self._处理菜谱(事务, 数据)
            事务.commit()
        except Exception as 异常:
            事务.rollback()
            self.日志.error(f"\n批量处理失败: {str(异常)}")

            # 批量失败后转为单条处理
            for 数据 in 批次:
                try:
                    self._处理菜谱(self.图数据库, 数据)
                except Exception as 异常:
                    self.日志.error(f"\n处理菜谱失败: {数据.get('name', '无名菜谱')} - {str(异常)}")

    def _处理菜谱(self, 事务, 数据: Dict[str, Any]):
        """处理单个菜谱数据

        Args:
            事务: 事务对象
            数据: 菜谱数据
        """
        # 创建菜谱节点
        菜谱 = Node(
            "菜谱",
            名称=数据["name"].strip(),
            图片链接=数据.get("img_url", "").strip()
        )
        事务.merge(菜谱, "菜谱", "名称")

        # 处理功效
        for 功效 in 数据.get("effects", []):
            if 功效.strip():
                功效节点 = Node("功效", 名称=功效.strip())
                事务.merge(功效节点, "功效", "名称")
                事务.create(Relationship(菜谱, "具有功效", 功效节点))

        # 处理食材
        for 食材 in 数据.get("ingredients", []):
            if isinstance(食材, str):
                if ":" in 食材:
                    名称, 用量 = 食材.split(":", 1)
                else:
                    名称, 用量 = 食材, "适量"
            else:
                名称, 用量 = str(食材), "适量"

            if 名称.strip():
                # 1. 先创建节点并设置名称（主属性）
                食材节点 = Node("食材", 名称=名称.strip())
                事务.merge(食材节点, "食材", "名称")

                # 2. 再补充设置用量属性
                食材节点["用量"] = 用量.strip()
                事务.push(食材节点)

                # 3. 创建关系
                事务.create(Relationship(菜谱, "使用食材", 食材节点))

                # 处理步骤
        for 序号, 步骤 in enumerate(数据.get("steps", []), 1):
            if 步骤.strip():
                步骤节点 = Node(
                    "步骤",
                    描述=步骤.strip(),
                    序号=序号
                )
                事务.create(步骤节点)
                事务.create(Relationship(菜谱, "包含步骤", 步骤节点))

        # 处理适宜人群
        for 人群 in 数据.get("suitable_group", []):
            if 人群.strip():
                人群节点 = Node("适宜人群", 描述=人群.strip())
                事务.merge(人群节点, "适宜人群", "描述")
                事务.create(Relationship(菜谱, "适宜于", 人群节点))

        # 处理禁忌人群
        for 人群 in 数据.get("taboo_group", []):
            if 人群.strip():
                人群节点 = Node("禁忌人群", 描述=人群.strip())
                事务.merge(人群节点, "禁忌人群", "描述")
                事务.create(Relationship(菜谱, "禁忌于", 人群节点))


if __name__ == "__main__":
    try:
        配置 = {
            "neo4j_uri": "bolt://localhost:7687",
            "neo4j_user": "neo4j",
            "neo4j_password": "12345678",
            "data_file": os.path.join(os.path.dirname(__file__), "caipu.jsonl")
        }

        if not os.path.exists(配置["data_file"]):
            raise FileNotFoundError(f"菜谱数据文件不存在: {配置['data_file']}")

        导入器 = Neo4j菜谱导入器(
            配置["neo4j_uri"],
            配置["neo4j_user"],
            配置["neo4j_password"],
            批量大小=200
        )
        导入器.导入JSON数据(配置["data_file"])

    except Exception as 异常:
        logging.error(f"程序异常: {str(异常)}", exc_info=True)
        exit(1)